package com.kwin.demo.api.util;

import org.apache.commons.codec.binary.Base64;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.PBEKeySpec;
import javax.crypto.spec.SecretKeySpec;
import java.nio.charset.StandardCharsets;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.KeySpec;

/**
 * @author kwin
 * @Date 2022/7/13 16:20
 **/
public class AESUtil {
    private static final String KEY_ALGORITHM = "AES";
    private static final String CIPHER_ALGORITHM = "AES/GCM/NoPadding";
    private static final int TAG_LENGTH_BIT = 128;
    private static final int IV_LENGTH_BYTE = 12;
    private static final int SALT_LENGTH_BYTE = 16;
    private static final int AES_KEY_BIT = 256;

    /**
     * 安全随机数生成
     * @param numBytes
     * @return
     */
    public static byte[] getRandomNonce(int numBytes) {
        byte[] nonce = new byte[numBytes];
        new SecureRandom().nextBytes(nonce);
        return nonce;
    }

    /**
     * 生成盐值
     * @return
     */
    public static String generateSalt() {
        byte[] salt = getRandomNonce(SALT_LENGTH_BYTE);
        return Base64.encodeBase64String(salt);
    }

    /**
     * 生成初始向量
     * @return
     */
    public static String generateIv() {
        byte[] iv = getRandomNonce(IV_LENGTH_BYTE);
        return Base64.encodeBase64String(iv);
    }

    /**
     * 生成256位秘钥
     * @return
     * @throws NoSuchAlgorithmException
     */
    public static SecretKey getAESKey() throws NoSuchAlgorithmException {
        KeyGenerator keyGen = KeyGenerator.getInstance(KEY_ALGORITHM);
        keyGen.init(AES_KEY_BIT, SecureRandom.getInstanceStrong());
        return keyGen.generateKey();
    }

    /**
     * 生成256位秘钥，base64编码
     * @return
     * @throws NoSuchAlgorithmException
     */
    public static String generateAESKey() throws NoSuchAlgorithmException {
        return Base64.encodeBase64String(getAESKey().getEncoded());
    }

    /**
     * 通过密码和盐值生成256位aes秘钥
     * @param password 密码
     * @param salt 盐值
     * @return
     * @throws NoSuchAlgorithmException
     * @throws InvalidKeySpecException
     */
    public static SecretKey getAESKeyFromPassword(char[] password, byte[] salt)
        throws NoSuchAlgorithmException, InvalidKeySpecException {

        SecretKeyFactory factory = SecretKeyFactory.getInstance("PBKDF2WithHmacSHA256");
        // iterationCount = 65536
        KeySpec spec = new PBEKeySpec(password, salt, 65536, AES_KEY_BIT);
        SecretKey secret = new SecretKeySpec(factory.generateSecret(spec).getEncoded(), KEY_ALGORITHM);
        return secret;

    }

    /**
     * 加密
     * @param source 明文
     * @param secret 秘钥
     * @param iv 初始向量
     * @return
     * @throws Exception
     */
    public static byte[] encrypt(byte[] source, SecretKey secret, byte[] iv) throws Exception {
        Cipher cipher = Cipher.getInstance(CIPHER_ALGORITHM);
        cipher.init(Cipher.ENCRYPT_MODE, secret, new GCMParameterSpec(TAG_LENGTH_BIT, iv));
        byte[] encryptedText = cipher.doFinal(source);
        return encryptedText;
    }

    /**
     * 加密，base64编码
     * @param source 明文
     * @param key 秘钥
     * @param iv 初始向量
     * @return
     * @throws Exception
     */
    public static String encrypt(String source, String key, String iv) throws Exception {
        SecretKey secretKey = new SecretKeySpec(Base64.decodeBase64(key), KEY_ALGORITHM);//getAESKeyFromPassword(password.toCharArray(), Base64.decodeBase64(salt));
        byte[] encryptedText = encrypt(source.getBytes(StandardCharsets.UTF_8), secretKey, Base64.decodeBase64(iv));
        return Base64.encodeBase64String(encryptedText);
    }

    /**
     * 加密，通过base64编码，秘钥通过password和盐值生成
     * @param source 明文
     * @param password 密码
     * @param iv 初始向量
     * @param salt 盐值
     * @return
     * @throws Exception
     */
    public static String encrypt(String source, String password, String iv, String salt) throws Exception {
        SecretKey secretKey = getAESKeyFromPassword(password.toCharArray(), Base64.decodeBase64(salt));
        byte[] encryptedText = encrypt(source.getBytes(StandardCharsets.UTF_8), secretKey, Base64.decodeBase64(iv));
        return Base64.encodeBase64String(encryptedText);
    }

    /**
     * 解密
     * @param encryptBytes 加密数组
     * @param secret 秘钥
     * @param iv 初始向量
     * @return
     * @throws Exception
     */
    public static String decrypt(byte[] encryptBytes, SecretKey secret, byte[] iv) throws Exception {

        Cipher cipher = Cipher.getInstance(CIPHER_ALGORITHM);
        cipher.init(Cipher.DECRYPT_MODE, secret, new GCMParameterSpec(TAG_LENGTH_BIT, iv));
        byte[] plainText = cipher.doFinal(encryptBytes);
        return new String(plainText, StandardCharsets.UTF_8);
    }

    /**
     * 解密
     * @param encryptStr 加密字符串
     * @param password 密码
     * @param iv 初始向量
     * @param salt 盐值
     * @return
     * @throws Exception
     */
    public static String decrypt(String encryptStr, String password, String iv, String salt) throws Exception {
        SecretKey secretKey = getAESKeyFromPassword(password.toCharArray(), Base64.decodeBase64(salt));
        return decrypt(Base64.decodeBase64(encryptStr), secretKey, Base64.decodeBase64(iv));
    }

    /**
     * 解密
     * @param encryptStr 加密字符串
     * @param key 秘钥
     * @param iv 初始向量
     * @return
     * @throws Exception
     */
    public static String decrypt(String encryptStr, String key, String iv) throws Exception {
        SecretKey secretKey = new SecretKeySpec(Base64.decodeBase64(key), KEY_ALGORITHM);
        return decrypt(Base64.decodeBase64(encryptStr), secretKey, Base64.decodeBase64(iv));
    }

}
